from netmiko import ConnectHandler
import os
import time
import threading
import queue
from AES import prpcrypt

pc = prpcrypt('*************')  # 初始化密钥
DEVICES_INFO_PATH = 'DEVICES.cfg'
# LOG_PATH = 'LOG'
THREADING_NUM = 10
# The number of threading
CMD_DELAY = 1
# sleep time
Q = queue.Queue(THREADING_NUM)
# FIFO
queueLock = threading.Lock()


# for print
def show_info(msg):
    queueLock.acquire()
    print(msg)
    queueLock.release()


####################################################################
def get_devices_info():
    # FORMAT:devicename ip username password cmd
    try:
        with open(DEVICES_INFO_PATH, 'r') as f:
            for line in f:
                if len(line) == 1:
                    # ignore blank line (len('\n')==1)
                    continue
                info_line = line.split()
                # using whitespace as the delimiter string(leading and trailing whitespace removed)
                if info_line[0][0] == '#':
                    # ignore comment line
                    continue
                else:
                    info_dict = {
                        'devicename': info_line[0],
                        'ip': info_line[1],
                        'username': pc.decrypt(info_line[2]),
                        'password': pc.decrypt(info_line[3]),
                        'cmd': info_line[4]
                    }
                    yield info_dict
    except FileNotFoundError as e:
        show_info('Can not find "{}"'.format(DEVICES_INFO_PATH))
    except IndexError as e:
        show_info('"{}" format error'.format(DEVICES_INFO_PATH))

    ####################################################################


def get_cmd_info(CMD_Info_Path):
    CMD = []
    DEVICES_TYPE = ''
    try:
        with open(CMD_Info_Path, 'r') as f:
            for line in f:
                if len(line) == 1:
                    continue
                # ignore blank line (len('\n')==1)
                cmd_line = line.strip()
                # leading and trailing whitespace removed
                if cmd_line[0] == '#':
                    continue
                # ignore comment line
                if cmd_line[:12] == 'device_type:':
                    DEVICES_TYPE = cmd_line[12:]
                    # get device_type
                else:
                    CMD.append(cmd_line)
            if DEVICES_TYPE == '':
                raise UnboundLocalError('No Device Type')
            return DEVICES_TYPE, CMD
    except FileNotFoundError as e:
        show_info('Can not find "{}"'.format(CMD_Info_Path))
    except UnboundLocalError as e:
        show_info('Can not find "device_type" in "{}"'.format(CMD_Info_Path))

    ####################################################################


def run_cmd(host, commands, device_type):
    Q.put(host)
    print('Start......{}-{}'.format(host['ip'], host['devicename']))
    try:
        host_netmiko = {
            'device_type': device_type,
            'ip': host['ip'],
            'username': host['username'],
            'password': host['password'],
            'port': 22
        }
        # Use netmiko:
        net_connect = ConnectHandler(**host_netmiko)

        # 执行指令
        net_connect.send_command(commands[0] + " " + host['devicename'] + ".cfg")
        time.sleep(CMD_DELAY)

        show_info('Finished...{}-{}'.format(host['ip'], host['devicename']))
        net_connect.disconnect()
    except:
        show_info('Failed.....{}-{} <Please check the fail.txt>'.format(host['ip'], host['devicename']))
    finally:
        Q.get()
        Q.task_done()

    ####################################################################


if __name__ == '__main__':
    pc = prpcrypt('*************')  # 初始化密钥
    LogTime = time.strftime('%Y-%m-%d_%H-%M-%S')
    threads = []
    if os.path.exists('LOG') == False:
        os.makedirs('LOG')
    # Create a log directory
    if os.path.exists('fail.txt') == True:
        os.remove('fail.txt')
    # remove 'fail.txt'

    hosts = get_devices_info()
    # get hosts
    try:
        for i in hosts:
            device_type, commands = get_cmd_info(i['cmd'])
            # get device_type and cmds
            task = threading.Thread(target=run_cmd, args=(i, commands, device_type))
            threads.append(task)
        for t in threads:
            t.start()
        for t in threads:
            t.join()
        show_info('All done.')
    except TypeError as e:
        show_info('Please check the error.')